from typing import Dict, Type
from app.core.storage.base import BaseStorage
from app.core.storage.local import LocalStorage
from app.core.storage.s3 import S3Storage
from app.core.storage.oss import OSSStorage
from app.core.storage.cos import COSStorage
from app.models.file_model import StorageProviderModel
from app.core.constants import StorageProviderEnum


class StorageFactory:
    """存储策略工厂"""

    # 存储策略映射
    _strategies: Dict[StorageProviderEnum, Type[BaseStorage]] = {
        StorageProviderEnum.LOCAL: LocalStorage,
        StorageProviderEnum.AWS_S3: S3Storage,
        StorageProviderEnum.ALIYUN_OSS: OSSStorage,
        StorageProviderEnum.TENCENT_COS: COSStorage,
    }

    @classmethod
    async def get_storage(cls, provider_id: int = None) -> BaseStorage:
        """
        获取存储策略

        Args:
            provider_id: 供应商ID，为None时使用默认供应商

        Returns:
            BaseStorage: 存储策略实例
        """
        # 获取供应商信息
        if provider_id:
            provider = await StorageProviderModel.get_or_none(id=provider_id, is_active=True)
        else:
            provider = await StorageProviderModel.get_or_none(is_default=True, is_active=True)

        if not provider:
            raise Exception("存储供应商不存在或未启用")

        # 获取对应的存储策略类
        strategy_cls = cls._strategies.get(provider.provider_type)
        if not strategy_cls:
            raise Exception("不支持的存储类型")

        # 创建存储策略实例
        return strategy_cls(
            access_key=provider.access_key,
            secret_key=provider.secret_key,
            endpoint=provider.endpoint,
            bucket=provider.bucket,
            region=provider.region,
            base_url=provider.base_url,
        )

    @classmethod
    def register_strategy(cls, provider_type: StorageProviderEnum, strategy_cls: Type[BaseStorage]):
        """
        注册新的存储策略

        Args:
            provider_type: 供应商类型
            strategy_cls: 策略类
        """
        cls._strategies[provider_type] = strategy_cls
